# coding: utf-8

import argparse
import yaml
import os, sys
import shutil
from pathlib import Path
import glob


def read_yaml(yaml_file):
    with open(yaml_file, "r", encoding="utf-8") as f:
        data = yaml.load(f.read(), yaml.FullLoader)
    return data


class X265Image:
    def __init__(self,
                 image_folder,
                 temp_folder: None):
        self.image_folder = image_folder
        # self.output_folder = output_folder
        self.temp_folder = temp_folder if temp_folder is not None else "./mkv_folder"   # mkv_folder
        # os.makedirs(self.output_folder, exist_ok=True)
        os.makedirs(self.temp_folder, exist_ok=True)
        self.ffmpeg = "ffmpeg"
        pass

    def encode(self, qp, suffix: str=".png"):
        image_paths = glob.glob(os.path.join(self.image_folder, "*"+suffix), recursive=True)
        for image in image_paths:
            print("Encoding: ", image)
            temp_name = os.path.basename(image).replace(suffix, ".mkv")
            # output_name = os.path.basename(image).replace(suffix, "_dec"+suffix)
            temp_path = os.path.join(self.temp_folder, temp_name)
            # output_path = os.path.join(self.output_folder, output_name)
            command1 = f"{self.ffmpeg} -i {image} -c:v libx265 -x265-params crf={qp} {temp_path}"
            os.system(command1)
            # 解码 mkv -> png
            # command2 = f"{self.ffmpeg} -i {temp_path} -r 1 -f image2 {output_path}"
            # os.system(command2)


def main():
    parse = argparse.ArgumentParser()
    parse.add_argument(
        "--yaml_path", "-yp", required=True, type=str,
        help="The yaml path of config file."
    )
    parse = parse.parse_args()
    config = read_yaml(parse.yaml_path)
    for crf in config["crf"]:
        # output_folder = os.path.join(config["save_foler"], str(crf))
        temp_folder = os.path.join(config["mkv_save_folder"], str(crf))
        x265 = X265Image(config["image_folder"], temp_folder)
        x265.encode(crf, suffix=".png")
    print("Complete all images.")


if __name__ == '__main__':
    main()
